"""ResNet."""
import mindspore.nn as nn
from mindspore.ops import operations as P

class ConvBNReLU(nn.Cell):
    def __init__(self, in_planes, out_planes, kernel_size=3, stride=1, groups=1):
        super(ConvBNReLU,self).__init__()
        padding = (kernel_size - 1)//2
        conv = nn.Conv2dBnAct(in_planes,out_planes,kernel_size,stride,pad_mode='pad',padding=padding,
                              groups=groups,has_bn=True,activation='relu')
        self.features = conv
    def construct(self,x):
        output = self.features(x)
        return output

class ResidualBlock(nn.Cell):
    expansion = 4
    def __init__(self,
                 in_channel,
                 out_channel,
                 strde=1):
        super(ResidualBlock,self).__init__()

        channel = out_channel // self.expansion
        self.conv1 = ConvBNReLU(in_channel,channel,kernel_size=1,stride=1)
        self.conv2 = ConvBNReLU(channel,channel,kernel_size=3,stride=strde)
        self.conv3 = nn.Conv2dBnAct(channel,out_channel,kernel_size=1,stride=1,
                                    pad_mode="same",padding=0,has_bias=True,activation='relu')

        self.down_sample = False
        if strde !=1 or in_channel != out_channel:
            self.down_sample = True
        self.down_sample_layer = None

        if self.down_sample:
            self.down_sample_layer = nn.Conv2dBnAct(in_channel, out_channel,
                                                    kernel_size=1, stride=stride,
                                                    pad_mode='same', padding=0, has_bn=True, activation='relu')
        self.add = P.TensorAdd()
        self.relu = P.ReLU()

    def construct(self,x):
        identity = x
        out = self.conv1(x)
        out = self.conv2(x)
        out = self.conv3(x)

        if self.down_sample:
            identity = self.down_sample_layer(identity)
        out = self.add(out,identity)
        out = self.relu(out)

        return out

class ResNet50(nn.Cell):
    """
    >>> ResNet(ResidualBlock,
        >>>        [3, 4, 6, 3],
        >>>        [64, 256, 512, 1024],
        >>>        [256, 512, 1024, 2048],
        >>>        [1, 2, 2, 2],
        >>>        10)
    """
    def __init__(self,
                 block,
                 layer_nums,
                 in_channels,
                 out_channels,
                 strides,
                 num_classes):
        super(ResNet50,self).__init__()

        if not len(layer_nums) == len(in_channels) ==len(out_channels) ==4:
            raise ValueError("the length of layer_num, in_channels, out_channels list must 4")

        self.conv1 = self._make_layer(block,
                                      layer_nums[0],
                                      in_channel=in_channels[0],
                                      out_channels=out_channels[0],
                                      stride=strides[0])
        self.maxpool = nn.MaxPool2d(kernel_size=3,stride=2,pad_mode='same')

        self.layer1 = self._make_layer(block,
                                       layer_nums[0],
                                       in_channel=in_channels[0],
                                       out_channels=out_channels[0],
                                       stride=strides[0])
        self.layer2 = self._make_layer(block,
                                       layer_nums[1],
                                       in_channel=in_channels,
                                       out_channel=out_channels,
                                       strides=strides[1])
        self.layer3 = self._make_layer(block,
                                       layer_nums[2],
                                       in_channel=in_channels,
                                       out_channel=out_channels,
                                       strides=strides[1])
        self.layer4 = self._make_layer(block,
                                       layer_nums[3],
                                       in_channel=in_channels,
                                       out_channel=out_channels,
                                       strides=strides[1])

        self.mean = P.ReduceMean(keep_dims=True)
        self.flatten = nn.Flatten()
        self.end_point = nn.DenseBnAct(out_channels[3],num_classes,has_bias=True,has_bn=False)

    def _make_layer(self,block,layer_num,in_channel,out_channel,stride):
        """
        >>> _make_layer(ResidualBlock, 3, 128, 256, 2)
        :param block: Resnet block.
        :param layer_num: Layer number.
        :param in_channel: Input channel.
        :param out_channel: Output channel.
        :param stride: Stride size for the first convolutional layer.
        :return: SequentialCell, the output layer.
        """
        layers = []
        resnet_block = block(in_channel, out_channel, stride=stride)
        layers.append(resnet_block)

        for _ in range(1,layer_num):
            resnet_block = block(out_channel,out_channel,stride=1)
            layers.append(resnet_block)

        return nn.SequentialCell(layers)

    def construct(self,x):
        x = self.conv1(x)
        c1 = self.maxpool(x)

        c2 = self.layer1(c1)
        c3 = self.layer2(c2)
        c4 = self.layer3(c3)
        c5 = self.layer4(c4)

        out = self.mean(c5,(2,3))
        out = self.flatten(out)
        out = self.end_point(out)
        return out

    def resnet50_quant(class_num=10):
        return ResNet(ResidualBlock,
                      [3, 4, 6, 3],
                      [64, 256, 512, 1024],
                      [256, 512, 1024, 2048],
                      [1, 2, 2, 2],
                      class_num)